import numpy as np

class DiffCalculator():
    def __init__(self,bit_num,target_bit):
        self.bit_num=bit_num
        self.target_bit=target_bit
        self.guess_key_diff=np.zeros((256))
    def _cal(self,traces, guess_midvalue,guess_key):
        index=np.arange(traces.shape[0])
        bit_zero=guess_midvalue==0
        traces0 = traces[index[bit_zero]]
        traces1 = traces[index[~bit_zero]]
        if len(traces0) == 0 or len(traces1) == 0:
            raise RuntimeError("Need to set a larger 'attack_steps' parameter!")
        traces_diff = np.mean(traces0, axis=0) - np.mean(traces1, axis=0)
        self.guess_key_diff[guess_key] += np.max(np.abs(traces_diff))
    def update(self,traces, guess_midvalue):
        for guess_key in range(256):
            if self.target_bit is None:
                for bit in range(self.bit_num):
                    self._cal(traces,guess_midvalue[:,guess_key,bit],guess_key)
            else:
                self._cal(traces, guess_midvalue[:, guess_key, self.target_bit], guess_key)
    def get_diff(self):
        return self.guess_key_diff

def attack(self):
    midvalue_shape = self.midvalue_attack.shape
    if len(midvalue_shape)!=3:
        raise ValueError("Please select right leakage model for attack('***MBit').")
    trs_num = self.traces_attack.shape[0]
    attack_steps = self.param_dict["attack_steps"]
    index = [i for i in range(0, trs_num, attack_steps)]
    index.append(trs_num)
    diff_calculator = DiffCalculator(midvalue_shape[2],self.target_bit)
    for i in range(len(index) - 1):
        diff_calculator.update(self.traces_attack[index[i]:index[i + 1]], self.midvalue_attack[index[i]:index[i + 1]])
        diff = diff_calculator.get_diff()
        yield diff, index[i + 1]

def get_guess_key(self, diff):
    return np.argmax(diff)

def get_rank(self,diff,real_key):
    corr_sort=np.argsort(np.argsort(-diff))
    return corr_sort[real_key]
